import numpy as np
import tensorflow as tf

from src.ga.interpreter import str2tensor


def get_tf_output(tensor, operator, variable=None):
    real_tensor = str2tensor(tensor)
    if operator == 'avg_pool':
        return tf.nn.avg_pool(real_tensor, 2, 2, 'SAME')
    elif operator == 'avg_pool3d':
        return tf.nn.avg_pool3d(real_tensor, 2, 2, 'SAME')
    elif operator == 'bias_add':
        # bias = tf.Variable(tf.random.normal([real_tensor.shape[-1]], stddev=0.01))
        return tf.nn.bias_add(real_tensor, variable)
    elif operator == 'conv1d':
        # my_filter = tf.Variable(tf.random.normal([1, real_tensor.shape[-1], 1], stddev=0.01))
        return tf.nn.conv1d(real_tensor, variable, stride=1, padding='SAME')
    elif operator == 'conv2d':
        # my_filter = tf.Variable(tf.random.normal([3, 3, real_tensor.shape[-1], 32], stddev=0.01))
        return tf.nn.conv2d(real_tensor, variable, strides=[1, 1, 1, 1], padding='SAME')
    elif operator == 'conv3d':
        # my_filter = tf.Variable(tf.random.normal([3, 3, 3, real_tensor.shape[-1], 32], stddev=0.01))
        return tf.nn.conv3d(real_tensor, variable, strides=[1, 1, 1, 1, 1], padding='SAME')
    elif operator == 'dilation2d':
        # my_filter = tf.Variable(tf.random.normal([3, 3, real_tensor.shape[-1]], stddev=0.01))
        return tf.nn.dilation2d(real_tensor, variable, strides=[1, 1, 1, 1], padding='SAME',
                                data_format='NHWC', dilations=[1, 1, 1, 1])
    elif operator == 'depthwise_conv2d':
        # my_filter = tf.Variable(tf.random.normal([3, 3, real_tensor.shape[-1], 32], stddev=0.01))
        return tf.nn.depthwise_conv2d(real_tensor, variable, strides=[1, 1, 1, 1], padding='SAME')
    elif operator == 'softmax':
        return tf.nn.softmax(real_tensor)
    elif operator == 'erosion2d':
        # my_filter = tf.Variable(tf.random.normal([3, 3, real_tensor.shape[-1]], stddev=0.01))
        return tf.nn.depthwise_conv2d(real_tensor, variable, strides=[1, 1, 1, 1], padding='SAME',
                                      data_format='NHWC', dilations=[1, 1, 1, 1])
    elif operator == 'log_softmax':
        return tf.nn.log_softmax(real_tensor)
    else:
        return ''
